# -*- coding: utf-8 -*-
"""
Created on Sun Jun 24 15:32:21 2018

@author: Luther
"""

from aip import AipNlp
import sim
import os
import time
import pandas as pd
from retrying import retry

""" 你的 APPID AK SK """
'''
APP_ID = '10771642'
API_KEY = '7QKEGf47U1trhYptTB3bM4gP'
SECRET_KEY = 'sdDSGirKljr4ijyRGYFOgP4VEtoP7W6a'
'''
APP_ID = '11788538' 
API_KEY = 'BcIhkEUdtaKmaqGtEIAXG3r2' 
SECRET_KEY = '6Rsdrv6C2FTqEVbEeXroIEWtGjOk3BSI' 

client = AipNlp(APP_ID, API_KEY, SECRET_KEY)

@retry(wait_fixed=0.2,stop_max_attempt_number=3)
def api(text):
    j = client.lexer(text[:10000])
    return j


def getne(text):
    ORG, PER, LOC, TIME = [],[],[],[]
    try:
        j = api(text)
    except:
        j = {}
    
    try:
        for i in j['items']:
            if i['ne'] == "ORG":
                ORG.append(i['item'])
            elif i['ne'] == "PER":
                PER.append(i['item'])
            elif i['ne'] == "LOC":
                LOC.append(i['item'])
            elif i['ne'] == "TIME":
                TIME.append(i['item'])
            else:
                pass
    except KeyError:
        ORG, PER, LOC, TIME = [],[],[],[]
    
    return {'ORG': ORG, 'PER': PER, 'LOC': LOC, 'TIME':TIME}


# getkw(title, content)
def compute(path, ret_path):
    org, per, loc, ttime = [],[],[],[]
    file_list = [
        x for x in os.listdir(path) if os.path.splitext(x)[-1] == '.txt'
    ]
    print("{:}个待检索文件读取成功!".format(len(file_list)))

    df = sim.getcode(file_list)
    #df = pd.read_csv(r'C:\users\luther\desktop\miss.csv')
    start_time = time.time()
    for i in range(len(df)):
        code = (6 - len(str(int(df.stkcd[i])))) * '0' + str(int(df.stkcd[i]))
        year = df.year[i]
        file = '{}-{}.txt'.format(year, code)
        text = sim.readtext(path, file)
        ret = getne(text)
        org.append(";".join(ret['ORG']))
        per.append(";".join(ret['PER']))
        loc.append(";".join(ret['LOC']))
        ttime.append(";".join(ret['TIME']))
        end_time = time.time()
        remain_time = (end_time - start_time) / (i + 1) * (len(df) - (i + 1))
        print(
            "\r{:>6.2f}% Done, Time Remained: {:04}:{:02}:{:02}".format(
                (i + 1) * 100 / len(df), int(remain_time // 3600),
                int(remain_time // 60 % 60), int(remain_time % 60)),
            end="")

    df['ORG'] = org
    df['PER'] = per
    df['LOC'] = loc
    df['TIME'] = ttime
    df.to_csv(ret_path, index=False)
